# -*- coding: utf-8 -*-  
'''
评价函数

@author: luoyi
Created on 2021年6月28日
'''
import tensorflow as tf

import utils.logger_factory as logf
import utils.conf as conf
import utils.relationships as rel
from models.tplinker.process import parse_y, relative_idx_et, relative_idx_sht


#    NER任务评价函数
class NERMetrics(tf.metrics.Metric):
    def __init__(self,
                 name='NERMetrics',
                 max_sen_len=conf.TPLINKER.get_max_sentence_len(),
                 fc_layer=None,
#                  input_layer=None,
                 **kwargs):
        super(NERMetrics, self).__init__(name=name, **kwargs)
        
        self._max_sen_len = max_sen_len
        
        self._fc_layer = fc_layer
#         self._input_layer = input_layer
        
        self._acc = self.add_weight('ner_metrics', initializer='zero', dtype=tf.float64)
        
        if (max_sen_len % 2 == 0): fold_num = int(((1 + max_sen_len-2) * (max_sen_len-2) / 2 + (max_sen_len-2) / 2) / (max_sen_len-2))
        else: fold_num = int((1 + max_sen_len-2) * (max_sen_len-2) / 2 / (max_sen_len-2))
        self._fold_num = fold_num
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''计算词性标注预测正确的数量
            TP = et_to_et预测正确的数量
            TF = et_no预测正确的数量
            acc = TP + TF / P
        '''
        #    取标注值    Tensor(batch_size, max_handshake_size, 2)，pad部分为-1
        Y_et_to_et, Y_et_no, \
            _, _, _, \
            _, _, _ = parse_y(y_true)
        #    取预测值
        X_et_to_et = self._fc_layer.get_et_to_et()
        
        #    取正样本数量，和索引
        P, idx_p = relative_idx_et(Y_et_to_et)
        #    取正样本预测正确的数量
        pred_p = tf.gather_nd(X_et_to_et, indices=idx_p)
        pred_p = tf.math.argmax(pred_p, axis=-1)
        TP = tf.math.count_nonzero(pred_p == 1, axis=-1)
        
        #    取负样本数量，和索引
        N, idx_n = relative_idx_et(Y_et_no)
        #    取正样本预测正确的数量
        pred_n = tf.gather_nd(X_et_to_et, indices=idx_n)
        pred_n = tf.math.argmax(pred_n, axis=-1)
        TN = tf.math.count_nonzero(pred_n == 1, axis=-1)
        
        #    计算准确率
        acc = tf.math.reduce_sum(TP + TN) / tf.math.reduce_sum(P + N)
        self._acc.assign(acc)
        
        tf.print('----------------------------------------------------------------', output_stream=logf.get_logger_filepath('tplinker_metrics'))
        tf.print('ner_metric acc:', acc, ' TP:', TP, ' TN:', TN,' P:', P, ' N:', N, output_stream=logf.get_logger_filepath('tplinker_metrics'))
        pass
    
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    pass


#    RE任务评价函数
class REMetrics(tf.metrics.Metric):
    def __init__(self,
                 name='REMetrics',
                 max_sen_len=conf.TPLINKER.get_max_sentence_len(),
                 fc_layer=None,
#                  input_layer=None,
                 rel_size=len(rel.id_rel),
                 **kwargs):
        super(REMetrics, self).__init__(name=name, **kwargs)
        
        self._max_sen_len = max_sen_len
        self._rel_size = rel_size
        
        self._fc_layer = fc_layer
#         self._input_layer = input_layer
        
        if (max_sen_len % 2 == 0): fold_num = int(((1 + max_sen_len-2) * (max_sen_len-2) / 2 + (max_sen_len-2) / 2) / (max_sen_len-2))
        else: fold_num = int((1 + max_sen_len-2) * (max_sen_len-2) / 2 / (max_sen_len-2))
        self._fold_num = fold_num
        
        self._acc = self.add_weight('re_metrics', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''关系握手准确率计算
            TP = sh_1 + sh_2 + st_1 + st_2
            TF = sh_no + st_no + rel_no
            acc = TP + TF / P
        '''
        X_sht_to_sht = self._fc_layer.get_sht_to_sht()          #    Tensor(batch_size, rel_size, ∑max_sen_len, 2, 3)
        X_sh = X_sht_to_sht[:, :, :, 0, :]                      #    Tensor(batch_size, rel_size, ∑max_sen_len, 3)
        X_st = X_sht_to_sht[:, :, :, 1, :]                      #    Tensor(batch_size, rel_size, ∑max_sen_len, 3)
        
        #    取标注值    Tensor(batch_size, max_handshake_size, 2)，pad部分为-1
        _, _, \
            Y_sh_to_sh_1, Y_sh_to_sh_2, Y_sh_to_sh_no, \
            Y_st_to_st_1, Y_st_to_st_2, Y_st_to_st_no = parse_y(y_true)
        
        #    sh相关正负样本数，及其索引
        num_sh_1, idx_sh_1 = relative_idx_sht(Y_sh_to_sh_1)
        num_sh_2, idx_sh_2 = relative_idx_sht(Y_sh_to_sh_2)
        num_sh_no, idx_sh_no = relative_idx_sht(Y_sh_to_sh_no)
        #    计算sh相关TP, TN, P, N
        pred_sh_1 = tf.gather_nd(X_sh, indices=idx_sh_1)
        pred_sh_1 = tf.math.argmax(pred_sh_1, axis=-1)
        TP_sh_1 = tf.math.count_nonzero(pred_sh_1 == 1)
        
        pred_sh_2 = tf.gather_nd(X_sh, indices=idx_sh_2)
        pred_sh_2 = tf.math.argmax(pred_sh_2, axis=-1)
        TP_sh_2 = tf.math.count_nonzero(pred_sh_2 == 2)
        
        pred_sh_no = tf.gather_nd(X_sh, indices=idx_sh_no)
        pred_sh_no = tf.math.argmax(pred_sh_no, axis=-1)
        TP_sh_no = tf.math.count_nonzero(pred_sh_no == 0)
        
        
        #    sh相关正负样本数，及其索引
        num_st_1, idx_st_1 = relative_idx_sht(Y_st_to_st_1)
        num_st_2, idx_st_2 = relative_idx_sht(Y_st_to_st_2)
        num_st_no, idx_st_no = relative_idx_sht(Y_st_to_st_no)
        #    计算sh相关TP, TN, P, N
        pred_st_1 = tf.gather_nd(X_st, indices=idx_st_1)
        pred_st_1 = tf.math.argmax(pred_st_1, axis=-1)
        TP_st_1 = tf.math.count_nonzero(pred_st_1 == 1)
        
        pred_st_2 = tf.gather_nd(X_st, indices=idx_st_2)
        pred_st_2 = tf.math.argmax(pred_st_2, axis=-1)
        TP_st_2 = tf.math.count_nonzero(pred_st_2 == 2)
        
        pred_st_no = tf.gather_nd(X_st, indices=idx_st_no)
        pred_st_no = tf.math.argmax(pred_st_no, axis=-1)
        TP_st_no = tf.math.count_nonzero(pred_st_no == 0)
        
        
        TP = TP_sh_1 + TP_sh_2 + TP_st_1 + TP_st_2
        P = num_sh_1 + num_sh_2 + num_st_1 + num_st_2
        TN = TP_sh_no + TP_st_no
        N = num_sh_no + num_st_no
        
        acc = tf.math.reduce_sum(TP + TN) / tf.math.reduce_sum(P + N)
        self._acc.assign(acc)
        
        tf.print('re_metric ', 'acc:', acc, ' TP:', TP, ' P:', P, ' TN:', TN, ' N:', N, output_stream=logf.get_logger_filepath('tplinker_metrics'))
        pass
    
    #    计算指定标注的数量
    def count_tagging(self, x, y, tag_val):
        '''
            @param x: Tensor(batch_size, rel_size, ∑max_sen_len, 3)
            @param y: Tensor(batch_size, max_handshake_size, 2): 
        '''
        idx = tf.where(y[:, :, 0] >= 0)                                     #    Tensor(sum_sh_1, 2)
        idx_B = idx[:, 0]
        idx = tf.gather_nd(y, indices=idx)
        idx = tf.concat([idx_B[:, tf.newaxis], idx], axis=-1)               #    Tensor(sum_sh_1, 3)    [idx_batc_size, idx_rel, idx_vec]
        pred = tf.gather_nd(x, indices=idx)                                 #    Tensor(sum_sh_1, 3)
        pred = tf.math.argmax(pred, axis=-1)
        count = tf.math.count_nonzero(pred == tag_val)
        return count
    
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    pass